#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/socket.h>
#include <linux/wireless.h>
#include <syslog.h>

/* MIMO Tx parameter, ShortGI, MCS, STBC, etc.  these are fields in TXWI. Don't change this definition!!! */
typedef union _MACHTTRANSMIT_SETTING {
	struct {
		unsigned short MCS:7;	/* MCS */
		unsigned short BW:1;	/*channel bandwidth 20MHz or 40 MHz */
		unsigned short ShortGI:1;
		unsigned short STBC:2;	/*SPACE */
		unsigned short rsv:3;
		unsigned short MODE:2;	/* Use definition MODE_xxx. */
	} field;
	unsigned short word;
} MACHTTRANSMIT_SETTING, *PMACHTTRANSMIT_SETTING;

typedef struct _RT_802_11_MAC_ENTRY {
	unsigned char ApIdx;
	unsigned char Addr[6];
	unsigned char Aid;
	unsigned char Psm;		/* 0:PWR_ACTIVE, 1:PWR_SAVE */
	unsigned char MimoPs;		/* 0:MMPS_STATIC, 1:MMPS_DYNAMIC, 3:MMPS_Enabled */
	char AvgRssi0;
	char AvgRssi1;
	char AvgRssi2;
	unsigned int ConnectedTime;
	MACHTTRANSMIT_SETTING TxRate;
	unsigned int LastRxRate;
	short StreamSnr[3];				/* BF SNR from RXWI. Units=0.25 dB. 22 dB offset removed */
	short SoundingRespSnr[3];			/* SNR from Sounding Response. Units=0.25 dB. 22 dB offset removed */
} RT_802_11_MAC_ENTRY, *PRT_802_11_MAC_ENTRY;

typedef struct _RT_802_11_MAC_TABLE {
	unsigned long Num;
	RT_802_11_MAC_ENTRY Entry[32];
} RT_802_11_MAC_TABLE, *PRT_802_11_MAC_TABLE;

struct survey_table
{
	char channel[4];
	char ssid[33];
	char bssid[20];
	char security[23];
	char *crypto;
};

RT_802_11_MAC_TABLE mt;
struct survey_table st[64];
int survey_count = 0;

#define RTPRIV_IOCTL_SET (SIOCIWFIRSTPRIV + 0x02)
void iwpriv(const char *name, const char *key, const char *val)
{
	int socket_id;
	struct iwreq wrq;
	char data[64];
	snprintf(data, 64, "%s=%s", key, val);
	socket_id = socket(AF_INET, SOCK_DGRAM, 0);
	strcpy(wrq.ifr_ifrn.ifrn_name, name);
	wrq.u.data.length = strlen(data);
	wrq.u.data.pointer = data;
	wrq.u.data.flags = 0;
	ioctl(socket_id, RTPRIV_IOCTL_SET, &wrq);
	close(socket_id);
}

static void next_field(char **line, char *output, int n) {
	char *l = *line;
	int i;

	memcpy(output, *line, n);
	*line = &l[n];

	for (i = n - 1; i > 0; i--) {
		if (output[i] != ' ')
			break;
		output[i] = '\0';
	}
}

#define RTPRIV_IOCTL_GSITESURVEY (SIOCIWFIRSTPRIV + 0x0D)
void wifi_site_survey(const char *ifname, const char* essid, int print)
{
	char *s = malloc(IW_SCAN_MAX_DATA);
	int ret;
	int socket_id;
	struct iwreq wrq;
	char *line, *start;
	iwpriv(ifname, "SiteSurvey", (essid ? essid : ""));
	sleep(5);
	memset(s, 0x00, IW_SCAN_MAX_DATA);
	strcpy(wrq.ifr_name, ifname);
	wrq.u.data.length = IW_SCAN_MAX_DATA;
	wrq.u.data.pointer = s;
	wrq.u.data.flags = 0;
	socket_id = socket(AF_INET, SOCK_DGRAM, 0);
	ret = ioctl(socket_id, RTPRIV_IOCTL_GSITESURVEY, &wrq);
	close(socket_id);
	if(ret != 0)
		goto out;
	if(wrq.u.data.length < 1)
		goto out;
	/* ioctl result starts with a newline, for some reason */
	start = s;
	while (*start == '\n')
		start++;
	line = strtok((char *)start, "\n");
	line = strtok(NULL, "\n");
	survey_count = 0;
	while(line && (survey_count < 64)) {
		next_field(&line, st[survey_count].channel, sizeof(st->channel));
		next_field(&line, st[survey_count].ssid, sizeof(st->ssid));
		next_field(&line, st[survey_count].bssid, sizeof(st->bssid));
		next_field(&line, st[survey_count].security, sizeof(st->security));
		line = strtok(NULL, "\n");
		st[survey_count].crypto = strstr(st[survey_count].security, "/");
		if (st[survey_count].crypto) {
			*st[survey_count].crypto = '\0';
			st[survey_count].crypto++;
		} else {
			st[survey_count].crypto = "";
		}
		survey_count++;
	}
	if (survey_count == 0 && !print)
		syslog(LOG_INFO, "No results");
out:
	free(s);
}

#define RTPRIV_IOCTL_GET_MAC_TABLE		(SIOCIWFIRSTPRIV + 0x0F)
static void get_assoclist(const char *ifname)
{
	int ret;
	int socket_id;
	struct iwreq wrq;
	
	strcpy(wrq.ifr_name, ifname);
	wrq.u.data.length = sizeof(mt);
	wrq.u.data.pointer = &mt;
	wrq.u.data.flags = 0;
	socket_id = socket(AF_INET, SOCK_DGRAM, 0);
	ret = ioctl(socket_id, RTPRIV_IOCTL_GET_MAC_TABLE, &wrq);
	close(socket_id);
	if(ret || wrq.u.data.length < 1)
		return;
}

int check_assoc(char *ifname)
{
	int socket_id, i;
	struct iwreq wrq;
	socket_id = socket(AF_INET, SOCK_DGRAM, 0);
	strcpy(wrq.ifr_ifrn.ifrn_name, ifname);
	ioctl(socket_id, SIOCGIWAP, &wrq);
	close(socket_id);
	for(i = 0; i < 6; i++)
		if(wrq.u.ap_addr.sa_data[i])
			return 1;
	return 0;
}

static void usage(const char *name)
{
	fprintf(stderr, "Usage: %s [OPTION]\n", name);
	fprintf(stderr, "	-d                                            Run in background\n");
	fprintf(stderr, "	-s <ifname>                                   scan\n");
	fprintf(stderr, "	-a <ifname>                                   assoclist\n");
	fprintf(stderr, "	-o <file>                                     write to file\n");
	fprintf(stderr, "	-q                                            quiet mode (don't print messages to stdout\n");
	fprintf(stderr, "	-j                                            Use JSON format to file\n");
	fprintf(stderr, "	-c <ifname>                                   Check whether or not associated\n");
}

int main(int argc, char **argv)
{
	int ch, i;
	int cmd = 0;
	int background = 0;
	int quiet = 0;
	int json = 0;
	const char *file = NULL;
	const char *ifname = NULL;
	
	FILE *fp = NULL;
	
	while ((ch = getopt(argc, argv, "s:a:c:r:e:k:o:dqj")) != -1) {
		switch (ch) {
		case 's':
			cmd = 's';
			ifname = optarg;
			break;
		
		case 'a':
			cmd = 'a';
			ifname = optarg;
			break;
			
		case 'c':
			cmd = 'c';
			ifname = optarg;
			break;
			
		case 'o':
			file = optarg;
			break;
			
		case 'q':
			quiet = 1;
			break;
			
		case 'd':
			background = 1;
			break;
		
		case 'j':
			json = 1;
			break;
		
		default:
			usage(argv[0]);
			exit(0);
		}
	}
	
	if (optind < argc || !cmd) {
		usage(argv[0]);
		exit(0);
	}

	if (background)
		daemon(0, 0);
	
	if (file) {
		fp = fopen(file , "w");
		if (!fp) {
			perror("fopen");
			exit(1);
		}
	}
	
	switch (cmd) {
	case 's':
		wifi_site_survey(ifname, NULL, 0);
		
		if (!quiet)
			printf("%-30s%-20s%-20s%-10s%-10s\n", "ssid", "bssid", "security", "crypto", "channel");
			
		if (fp) {
			if (json)
				fprintf(fp, "{\"list\":[");
			else
				fprintf(fp, "%-30s%-20s%-20s%-10s%-10s\n", "ssid", "bssid", "security", "crypto", "channel");
		}
		
		for(i = 0; i < survey_count; i++) {
			if (!quiet)
				printf("%-30s%-20s%-20s%-10s%-10s\n", st[i].ssid, st[i].bssid, st[i].security, st[i].crypto, st[i].channel);
			
			if (fp) {
				if (json) {
					fprintf(fp, "{\"ssid\":\"%s\", \"bssid\":\"%s\", \"security\":\"%s\", \"crypto\":\"%s\", \"channel\":\"%s\"}", 
						st[i].ssid, st[i].bssid, st[i].security, st[i].crypto, st[i].channel);
					
					if (i < survey_count - 1)
						fprintf(fp, ",");
				} else {
					fprintf(fp, "%-30s%-20s%-20s%-10s%-10s\n", st[i].ssid, st[i].bssid, st[i].security, st[i].crypto, st[i].channel);
				}
			}
		}
		
		if (fp && json)
		fprintf(fp, "]}");
		
		break;
		
	case 'a':
		get_assoclist(ifname);
		if (!quiet)
			printf("%-20s\n", "mac");	
		
		if (fp)
			fprintf(fp, "%-20s\n", "mac");	
		
		for (i = 0; i < mt.Num; i++) {
			RT_802_11_MAC_ENTRY *Entry;
			unsigned char *mac;
			
			Entry = mt.Entry + i;
			mac = Entry->Addr;
			
			if (!quiet)
				printf("%02X:%02X:%02X:%02X:%02X:%02X   \n", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
			
			if (fp)
				fprintf(fp, "%02X:%02X:%02X:%02X:%02X:%02X   \n", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
		}
		break;
	
	case 'c':
		printf("%s is %sassociated\n", ifname, check_assoc(ifname) ? "not " : "");
		break;
		
	default:
		usage(argv[0]);
		break;
	}
	
quit:	
	if (fp)
		fclose(fp);
	
	return 0;
}
